#!/usr/bin/env python

import sys
import os
import socket
import select
import subprocess
import logging

from console import *
from selector import *
from multiplexer import *


class dummyHandler:
    def __init__(self, tag):
        self.tag = tag
    def handle(self, cont):
        sys.stderr.write(self.tag + ": " + cont + "\n")
    def close(self):
        sys.stderr.write(self.tag + ": closed\n")



class i2oHandler (selector.sender):
    def __init__(self, outSock, selector, ch):
        super(i2oHandler, self).__init__(outSock, selector)
        self.ch = ch
        self.closeHook = self._closeHook

    def handle(self, cont):
        self.send(cont)

    def close(self):
        super(i2oHandler, self).close()

    def _closeHook(self):
#        sys.stderr.write("--" + self.ch.chId + ": i2oClosed\n")
        self.ch.i2oClosed()

    def __str__(self):
        return "[i2oHandler:" + self.ch.chId + "]"


class o2iForwarder(selector.reader):
    def __init__(self, inSock, sinkTerminal, ch):
        super(o2iForwarder, self).__init__(inSock, ch.tunnel.selector)
        self.sinkTerminal = sinkTerminal
        self.ch = ch

    def got(self, cont):
        self.sinkTerminal.write(cont)

    def closed(self):
        self.ch.o2iClosed()

    def __str__(self):
        return "[o2iForwarder:" + self.ch.chId + "]"

    def throttle(self, flag):
        self.reading = flag
        if flag:
            self.ch.tunnel.selector.addIn(self)

import fcntl

class drainer:
    """ just drain the inputs """
    def __init__(self, f, sel):
        self.f = f
        flag = fcntl.fcntl(self.f, fcntl.F_GETFL)
        flag = flag | os.O_NONBLOCK
        fcntl.fcntl(self.f, fcntl.F_SETFL, flag)
        sel.addIn(self)

    def fileno(self):
        return self.f.fileno()

    def handleIn(self):
        while True:
            try:
                buf = self.f.read(4096)
                if len(buf) == 0:   # this means peer shutted down
                    return False
                else:
                    logging.debug(buf)
                    continue
            except IOError, (num, mes):
                if num == errno.EAGAIN:
                    return True

    def __str__(self):
        return "[drainer]"


class lForwarder:
    def __init__(self, tuple, tunnel):
        (self.laddr, self.lport, self.remoteHost, self.remotePort) = tuple
        self.tunnel = tunnel

    def accept(self):
        self.acc = self.tunnel.selector.accept(self.laddr, self.lport, self.acceptCallback)
    
    def getPort(self):
        (_h, port) = self.acc.ss.getsockname() 
        return port
        
    def acceptCallback(self, success, val):
        """ connected """
        if success:
            s, addr = val
            s.setblocking(0)
            self.tunnel.accepted(s, self.accId)
        return True

    def setAccId(self, id):
        self.accId = id

    def stopAccept(self):
        self.tunnel.selector.removeIn(self.acc)
        self.acc.close()

class rForwarder:
    def __init__(self, tuple, tunnel):
        (self.acceptingAddr, self.acceptingPort, self.host, self.port) = tuple
        self.acceptPort = None
        pass

    def setRealPort(self, port):
        self.acceptPort = port

    def getPort(self):
        """ will be called from console. """
        return self.acceptPort

    def setAccId(self, id):
        self.accId = id

class channel:
    counter = 0
    def __init__(self, accId, sock):
        self.accId = accId
        self.sock = sock
        self.chId = "ch" + str(channel.counter)
        channel.counter += 1
        self.sink = None
        self.handler = None
        self.o2iForwarder = None
        self.tunnel = None
        self.closingCounter = 0

    def activate(self, tunnel):
        self.tunnel = tunnel
        selector = tunnel.selector
        multiplexer = tunnel.multiplexer
        self.sink = multiplexer.getTerminal(self.chId)
        self.o2iForwarder = o2iForwarder(self.sock, self.sink, self)
        self.handler = i2oHandler(self.sock, selector, self)
        self.handler.setThrottle(self.throttleSrc)
        multiplexer.registerHandler(self.chId, self.handler)

    def _checkShutdown(self):
       self.closingCounter += 1
       if self.closingCounter >= 2: 
            self.sock.close()
            self.tunnel.selector.removeIn(self.o2iForwarder)
            self.tunnel.channelClosed(self)

    def o2iClosed(self):
        self.sink.close()
        try:
            self.sock.shutdown(socket.SHUT_RD)
        except socket.error:
            pass
        self._checkShutdown()

    def i2oClosed(self):
        self._checkShutdown()
        
    def throttleDst(self, bool):
        self.o2iForwarder.throttle(bool)

    def throttleSrc(self, bool):
        self.tunnel.throttle(self.chId, bool)

    def forceClose(self):
        self.handler.close()
#        self.o2iForwarder.closed()

    def __str__(self):
        return "[channel:" + self.chId + "]"
        
class tunnel:
    REP = "REP"
    HELLO = "HELLO"
    START_ACCEPT = "START_ACCEPT"
    STOP_ACCEPT = "STOP_ACCEPT"

    INIT_FORWARD = "INIT_FORWARD"
    ACCEPTED = "ACCEPTED"
    THROTTLE = "THROTTLE"
    QUIT = "QUIT"
    
    def __init__(self, multiplexer, selector):
        self.multiplexer = multiplexer
        self.selector = selector
        self.peer = multiplexer.getTerminal("peer")
        multiplexer.registerHandler("peer", self)
        self.accCounter = 0
        self.forwarderMap = {}
        self.reqCounter = 0
        self.reqMap = {}
        self.chCounter = 0
        self.accDstMap = {}
        self.accChMap = {}
        self.chMap = {}
        
    def errln(self, *args):
        logging.debug(" ".join(*args))

    def accepted(self, inSock, accId):
        # allocate channelName, register to channel map
        ch = channel(accId, inSock)
        try:
            self.accChMap[accId].append(ch)
        except KeyError:
            self.accChMap[accId] = [ch]
         # send accepted message
        self._accepted(ch)

    def _accepted(self, ch):
        reqId = self._getReqId()
        self._sendMes(self.ACCEPTED, reqId, ch.accId, ch.chId)
        self.reqMap[reqId] = lambda x: self._acceptedRepHandler(ch, x)

    def _acceptedHandler(self, reqId, accId, chId):
        # try to connect
        target = self.accDstMap[accId]
        self.selector.connect(target,
                              lambda succes, val: self._acceptedHandler_connect_callback(reqId, accId, chId, succes, val)) 
                              

    def _acceptedHandler_connect_callback(self, reqId, accId, chId, success, val):        
        if success:
        # create channel
        # return reply
            s = val
            ch = channel(accId, s)
            try:    
                self.accChMap[accId].append(ch)
            except KeyError:
                self.accChMap[accId] = [ch]
            ch.activate(self)
            self.chMap[ch.chId] =  ch
            self._reqSuccess(reqId)
        else:
            no, msg = val
            self._reqFail(reqId, msg)
    
    def _acceptedRepHandler(self, ch, x):
        if x[0] == "S":
            ch.activate(self)
            self.chMap[ch.chId] =  ch
        else:
            self.errln("peer failed to handle 'accepted' ", x[1:])

    def throttle(self, chId, flag):
        self._sendMes(self.THROTTLE, "ON" if flag else "OFF")

    def _throttleHandler(self, chId, flagString):
        flag = (flagString == "ON")
        try:
            self.chMap[chId].throttleDst(flag)
        except KeyError:
            self.errln("in _throttleHandler, cannot find chId ", chId)
           
    def lForward(self, tuple, callback):
        """try to open the port""" 
        forwarder = lForwarder(tuple, self)
        try:
            forwarder.accept()
            callback((True, forwarder))
            self._initForward(forwarder)
        except socket.error, (no, msg):
            callback((False, msg))

    def _initForward(self, forwarder):
        accId = self._getAccId()
        reqId = self._getReqId()
        self.forwarderMap[accId] = forwarder
        forwarder.setAccId(accId)
        self._sendMes(self.INIT_FORWARD, reqId, accId, forwarder.remoteHost, forwarder.remotePort)            
        self.reqMap[reqId] = lambda x: self._initForwardRepHandler(forwarder, x)

    def _initForwardRepHandler(self, forwarder, x):
        if x[0] == "S":
            pass
        else:
            logging.debug(" ".join(X[1:]))

    def _initForwardHandle(self, accId, host, port):
        self.accDstMap[accId] = (host, port)
        return True
    
    def rForward(self, tuple, callback):
        forwarder = rForwarder(tuple, self)
        # send start_accept 
        addr = forwarder.acceptingAddr
        port = forwarder.acceptingPort
        accId = self._getAccId()
        reqId = self._getReqId()
        forwarder.setAccId(accId)
        self._sendMes(self.START_ACCEPT, reqId, accId, addr, port)
        self.reqMap[reqId] = lambda x: self._startAcceptRepHandler(forwarder, accId, callback, x)

    def _startAcceptHandle(self, accId, addr, port):
        _forwarder = lForwarder((addr, port, None, None), self)
        _forwarder.setAccId(accId)
        try:
            _forwarder.accept()
            # success
            self.forwarderMap[accId] = _forwarder
            return True, _forwarder.getPort()
        except socket.error, (no, msg):
            return False, msg

    def _startAcceptRepHandler(self, forwarder, accId, callback, x):
        if x[0] == "S":
            self.forwarderMap[accId] = forwarder
            self.accDstMap[accId] = (forwarder.host, forwarder.port)
            forwarder.setRealPort(x[1])
            callback((True, forwarder))
        else:
            callback((False, " ".join(x[1:])))

##############################################

    def stopAccept(self, forwarder, callback):
        "just stop to accept any farther incoming connect"
        if isinstance(forwarder, lForwarder):
            forwarder.stopAccept()
            callback((True,None))
        else:
            reqId = self._getReqId()            
            accId = forwarder.accId
            self._sendMes(self.STOP_ACCEPT, reqId, accId)
            self.reqMap[reqId] = lambda x: self._stopAcceptRepHandler(forwarder, accId, callback, x)
    
    def _stopAcceptHandle(self, accId):
        try:
            forwarder = self.forwarderMap[accId]
            forwarder.stopAccept()
            return True, "success"
        except KeyError:
            return False, "cannot find forwarder at remote node"

    def _stopAcceptRepHandler(self, forwarder, accId, callback, x):
        if x[0] == "S":
            callback((True, None))
        else:
            callback((False, " ".join(x[1:])))
        
##############################################

    def shutdown(self, forwarder, callback):
        "shutdown all the living connection"
        accId = forwarder.accId
        self.stopAccept(forwarder, lambda x: self.shutdownChannel(x, accId, callback))
        
        
    def shutdownChannel(self, (success, message), accId, callback):
        try:
            list = self.accChMap[accId]
            for ch in list:
                ch.forceClose()
            del self.accChMap[accId]
        except KeyError:
            #ignore
            pass
        callback((success, message))
        

    def hello(self):
        self._sendMes(self.HELLO)

    def handle(self, data):  
        """this will called when the data comes from peer.
           data should be a line terminated with a \n 
        """ 
        logging.debug(data)
        array = data.strip().split(" ")
        if array[0] == self.HELLO: 
#            sys.stderr.write("HELLO detected\n")
            pass
        elif array[0] == self.REP:
            reqId = array[1]
            try:
                repHandler = self.reqMap[reqId]
                repHandler(array[2:])
            except KeyError:
                logging.debug("no ID " + reqId + " \n")
        elif array[0] == self.INIT_FORWARD:
            reqId = array[1]
            if self._initForwardHandle(array[2], array[3], int(array[4])):
                self._reqSuccess(reqId)
            else:
                self._reqFail(reqId, "will not happen")
        elif array[0] == self.START_ACCEPT:      
            reqId = array[1]
            success, port = self._startAcceptHandle(array[2], array[3], array[4])
            if success:
                self._reqSuccess(reqId, port)
            else:
                mes = port
                self._reqFail(reqId, mes)
        elif array[0] == self.STOP_ACCEPT:      
            reqId = array[1]
            success, mes = self._stopAcceptHandle(array[2])
            if success:
                self._reqSuccess(reqId, mes)
            else:
                self._reqFail(reqId, mes)
        elif array[0] == self.ACCEPTED:
            reqId = array[1]
            self._acceptedHandler(reqId, array[2], array[3])
        elif array[0] == self.THROTTLE:
            self._throttleHandler(array[1], array[2])

        elif array[0] == self.QUIT:
            sys.exit(0)
        return True


#    def _handle_accepted(self, reqid, channel, target):
#        connector(target)
        
#        (success, info) = self._connect(target)
#        if success:
#            s = info
#            self._reqSuccess(reqid)
#            self.multiplexer.registerHandler(channel, i2oHandler(s))
#            o2iForwarder(s, self.multiplexer.getTerminal(channel))
#        else:
#            self._reqFail(reqid, info)
           
    def _reqSuccess(self, reqId, *args):
        self._sendMes(self.REP, reqId, "S", *args)

    def _reqFail(self, reqId, *args):
        self._sendMes(self.REP, reqId, "F", *args)

    def _sendMes(self, *args):
        self.peer.write(" ".join(str(x) for x in args) + "\n")
    
    def _connect(self, target):
        s = socket.socket()
        try:
            s.connect(target)
        except socket.error, (errno, errMsg):
            return (False, errMsg)
        return (True, s)
    
    def channelClosed(self, ch):
        try:
            del self.chMap[ch.chId]
        except KeyError:
            pass

    def close(self):
        pass

    def quit(self):
        self._sendMes(self.QUIT)
        sys.exit(0)
    
    def _getAccId(self):
        tmp = self.accCounter
        self.accCounter += 1
        return "acc" + str(tmp)
    
    def _getReqId(self):
        tmp = self.reqCounter
        self.reqCounter += 1
        return "req" + str(tmp)


def server(port):
    ss = socket.socket()
    ss.bind(("localhost", port))
    ss.listen(10)
    (s, addr) = ss.accept()
    f = s.makefile()
    _server0(f, f)

def client(port):
    s = socket.socket()
    s.connect(("localhost", port))
    f = s.makefile()
    _client0(f, f)

f_options = []
    
def _client0(fout, fin, ferr=None):
    sel = selector()
    fin.readline()
    mul = multiplexer(fout, fin)
    tun = tunnel(mul, sel)
    con = console(sys.stdin, sys.stdout, sys.stderr, tun)
    sel.addIn(con)
    sel.addIn(mul)
    if ferr:
        drainer(ferr, sel)
    tun.hello()
    for l in f_options:
        con.handleLine(l)
    sel.selectLoop()

def _server0(fout, fin):
    sel = selector()
    fout.write("hello\n")
    mul = multiplexer(fout, fin)
    tun = tunnel(mul, sel)
    sel.addIn(mul)
    tun.hello()
    sel.selectLoop()
    
def filter(generator):
    for line in generator:
        if len(line.strip()) != 0:
            if not (line.startswith("#") or line.startswith(" ") or line.startswith("\t")):
                yield("\n")
            yield line


ssh_options = []


def invoke(host, thisFile):
    p = subprocess.Popen(["ssh"] + ssh_options +[host, "python -i" ], 
              stdin = subprocess.PIPE, 
              stdout = subprocess.PIPE,
              stderr = subprocess.PIPE) 
    p.stdin.write("__name__ = 'server'\n")
    for line in filter(file(thisFile, "r")):
        p.stdin.write(line) 
    p.stdin.flush()
    return (p.stdin, p.stdout, p.stderr) 


def usage(val):
    sys.stderr.write("""USAGE: ssht [OPTIONS] host
OPTIONS: -ssh-options-start SSH_OPTIONS .. -ssh-options-end
         -R [bind_address]:port:host:hostport
         -L [bind_address]:port:host:hostport
""")
    sys.exit(val)

def main():
    if __name__ == "__main__":
#        logging.basicConfig(level=logging.DEBUG)
        global ssh_options, f_options
        mode ="ssh" 
        targetHost = None
        try:
            i = 1
            ssh_option_mode = False
            while i < len(sys.argv):
                if not ssh_option_mode:
                    if sys.argv[i] == "-R":
                        f_options.append('R ' + sys.argv[i+1])
                        i += 1
                    elif sys.argv[i] == "-L":
                        f_options.append('L ' + sys.argv[i+1])
                        i += 1
                    elif sys.argv[i] == "-ssh_options_begin":
                        ssh_option_mode = True
                    elif sys.argv[i] == "-server":
                        mode = "server"
                    elif sys.argv[i] == "-client":
                        mode = "client"
                    elif sys.argv[i].startswith('-'):
                        usage(1)
                    else:
                        if targetHost:
                            usage(1)
                        else:
                            targetHost = sys.argv[i]
                else:
                    if sys.argv[i] == "-ssh_options_end":
                        ssh_option_mode = False
                    else:
                        ssh_options.append(sys.argv[i])
                i += 1
        
        except IndexError:           
            usage(1)

        if mode == "server":
            server(int(sys.argv[2]))
        elif mode == "client":
            client(int(sys.argv[2]))
        elif mode == "ssh":
            if not targetHost:
                usage(1)
            fin, fout, ferr = invoke(targetHost, sys.argv[0]) 
            _client0(fin, fout, ferr=ferr)
    elif __name__ == "server":
        _server0(sys.stdout, sys.stdin)
        
main()
